#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
from scapy.all import *
from typing import Dict, List, Tuple
import logging

class FeatureExtractor:
    """VPN流量特征提取器"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
    def extract_packet_features(self, packet: Packet) -> Dict[str, float]:
        """
        从单个数据包中提取特征
        """
        features = {}
        
        # 1. 基本特征
        features.update(self._extract_basic_features(packet))
        
        # 2. 统计特征
        features.update(self._extract_statistical_features(packet))
        
        # 3. 协议特征
        features.update(self._extract_protocol_features(packet))
        
        # 4. 时序特征
        features.update(self._extract_temporal_features(packet))
        
        return features
    
    def _extract_basic_features(self, packet: Packet) -> Dict[str, float]:
        """提取基本特征"""
        features = {}
        
        # 包大小
        features['packet_size'] = len(packet)
        
        # IP层特征
        if packet.haslayer(IP):
            features['ip_ttl'] = packet[IP].ttl
            features['ip_proto'] = packet[IP].proto
            
        # TCP/UDP层特征
        if packet.haslayer(TCP):
            features['tcp_sport'] = packet[TCP].sport
            features['tcp_dport'] = packet[TCP].dport
            features['tcp_flags'] = packet[TCP].flags
        elif packet.haslayer(UDP):
            features['udp_sport'] = packet[UDP].sport
            features['udp_dport'] = packet[UDP].dport
            
        return features
    
    def _extract_statistical_features(self, packet: Packet) -> Dict[str, float]:
        """提取统计特征"""
        features = {}
        
        if packet.haslayer(Raw):
            raw_data = bytes(packet[Raw])
            
            # 熵值计算
            features['entropy'] = self._calculate_entropy(raw_data)
            
            # 字节分布特征
            features.update(self._calculate_byte_distribution(raw_data))
            
        return features
    
    def _extract_protocol_features(self, packet: Packet) -> Dict[str, float]:
        """提取协议特征"""
        features = {}
        
        # TLS/SSL特征
        if packet.haslayer(TLS):
            features['tls_version'] = packet[TLS].version
            features['tls_handshake_type'] = packet[TLS].type
            
        # GRE特征
        if packet.haslayer(GRE):
            features['gre_proto'] = packet[GRE].proto
            
        # ESP特征
        if packet.haslayer(ESP):
            features['esp_spi'] = packet[ESP].spi
            
        return features
    
    def _extract_temporal_features(self, packet: Packet) -> Dict[str, float]:
        """提取时序特征"""
        features = {}
        
        # 时间戳特征
        features['timestamp'] = packet.time
        
        return features
    
    def _calculate_entropy(self, data: bytes) -> float:
        """计算数据熵值"""
        if not data:
            return 0.0
            
        counts = np.zeros(256)
        for byte in data:
            counts[byte] += 1
            
        # 计算概率
        probs = counts / len(data)
        probs = probs[probs > 0]
        
        # 计算熵值
        entropy = -np.sum(probs * np.log2(probs))
        return float(entropy)
    
    def _calculate_byte_distribution(self, data: bytes) -> Dict[str, float]:
        """计算字节分布特征"""
        features = {}
        
        if not data:
            return features
            
        # 计算字节频率
        counts = np.zeros(256)
        for byte in data:
            counts[byte] += 1
            
        # 归一化
        probs = counts / len(data)
        
        # 提取统计特征
        features['byte_mean'] = float(np.mean(probs))
        features['byte_std'] = float(np.std(probs))
        features['byte_max'] = float(np.max(probs))
        features['byte_min'] = float(np.min(probs))
        
        return features 